#include "structMeshPartitioning.hpp"
#include "mpiH_Comm.hpp"
#include "structMeshPart.hpp"
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_string.hpp>
#include <cmath>

using namespace hms;

using Catch::Matchers::ContainsSubstring;

TEST_CASE( "Integer factorisation" ){
	REQUIRE( factors(1) == std::vector<int>{1} );
	REQUIRE( factors(2) == std::vector<int>{1,2} );
	REQUIRE( factors(3) == std::vector<int>{1,3} );
	REQUIRE( factors(4) == std::vector<int>{1,2,4} );
	REQUIRE( factors(5) == std::vector<int>{1,5} );
	REQUIRE( factors(6) == std::vector<int>{1,2,3,6} );
	REQUIRE( factors(7) == std::vector<int>{1,7} );
	REQUIRE( factors(8) == std::vector<int>{1,2,4,8} );
	REQUIRE( factors(9) == std::vector<int>{1,3,9} );
	REQUIRE( factors(10) == std::vector<int>{1,2,5,10} );
	REQUIRE( factors(11) == std::vector<int>{1,11} );
	REQUIRE( factors(12) == std::vector<int>{1,2,3,4,6,12} );
	REQUIRE( factors(13) == std::vector<int>{1,13} );
	REQUIRE( factors(14) == std::vector<int>{1,2,7,14} );
	REQUIRE( factors(15) == std::vector<int>{1,3,5,15} );
	REQUIRE( factors(16) == std::vector<int>{1,2,4,8,16} );
	REQUIRE( factors(17) == std::vector<int>{1,17} );
	REQUIRE( factors(18) == std::vector<int>{1,2,3,6,9,18} );
	REQUIRE( factors(19) == std::vector<int>{1,19} );
	REQUIRE( factors(20) == std::vector<int>{1,2,4,5,10,20} );
	REQUIRE( factors(21) == std::vector<int>{1,3,7,21} );
	REQUIRE( factors(22) == std::vector<int>{1,2,11,22} );
	REQUIRE( factors(23) == std::vector<int>{1,23} );
	REQUIRE( factors(24) == std::vector<int>{1,2,3,4,6,8,12,24} );
}

TEST_CASE( "roundToFactor: Rounding a float to nearest factor of an int" ){
	REQUIRE( roundToFactor( 1, factors(1) ) == 1 );
	REQUIRE( roundToFactor( 1, factors(6) ) == 1 );
	REQUIRE( roundToFactor( 2, factors(6) ) == 2 );
	REQUIRE( roundToFactor( 3, factors(6) ) == 3 );
	REQUIRE( roundToFactor( 4, factors(6) ) == 3 );
	REQUIRE( roundToFactor( 5, factors(6) ) == 6 );
	REQUIRE( roundToFactor( 6, factors(6) ) == 6 );
	REQUIRE( roundToFactor( std::sqrt(6), factors(6) ) == 2 );
}

TEST_CASE( "structMeshNParts: Determining optimal initial partition dimensions"
){
	/* first two parameters are the mesh size in x and y directions, second is
	 * the number of nodes/ranks. Output is the number of nodes in either
	 * dimension */
	REQUIRE( (structMeshNParts( 12, 12,  2 ) == Array2i{ 2, 1}).all() );
	REQUIRE( (structMeshNParts( 12, 12,  3 ) == Array2i{ 3, 1}).all() );
	REQUIRE( (structMeshNParts( 12, 12,  4 ) == Array2i{ 2, 2}).all() );
	REQUIRE( (structMeshNParts( 12, 12,  5 ) == Array2i{ 5, 1}).all() );
	REQUIRE( (structMeshNParts( 12, 12,  6 ) == Array2i{ 3, 2}).all() );
	REQUIRE( (structMeshNParts( 12, 12,  8 ) == Array2i{ 4, 2}).all() );
	REQUIRE( (structMeshNParts( 12, 12,  9 ) == Array2i{ 3, 3}).all() );
	REQUIRE( (structMeshNParts( 12, 12, 10 ) == Array2i{ 5, 2}).all() );
	REQUIRE( (structMeshNParts( 12, 12, 11 ) == Array2i{11, 1}).all() );
	REQUIRE( (structMeshNParts( 12, 12, 12 ) == Array2i{ 4, 3}).all() );
	REQUIRE_THROWS_WITH(
		structMeshNParts( 12, 12, 13 ), ContainsSubstring( "Too many nodes" )
	);
}

TEST_CASE( "structMeshPartSize: Determining optimal initial partition sizes" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	auto test = [&](
		Index nx_g, Index ny_g, /* global mesh sizes */
		Index a, Index b,       /* number of partitions in x and y */
		Index nx_l, Index ny_l  /* expected local mesh sizes */
	){
		/* this test only makes sense if there are enough nodes available */
		if ( a*b > comm.size() ) return;
		/* communicator must be split off, so that number of parts equals number
		 * of ranks in communicator */
		MPI_Comm_H split =
			comm.split( comm.rank() < (a*b) ? 0 : MPI_UNDEFINED, comm.rank() );
		if (split)
			REQUIRE( (
				structMeshPartSize(Array2i{nx_g, ny_g}, Array2i{a,b}, split) ==
				Array2i{nx_l, ny_l}
			).all() );
	};

	test( 12, 12,  1,  1, 12, 12 );
	test( 12, 12,  2,  1,  6, 12 );
	test( 12, 12,  3,  1,  4, 12 );
	test( 12, 12,  2,  2,  6,  6 );
	test( 12, 12,  5,  1,  2+((comm.rank()+1) % 5==0?2:0), 12 );
	test( 12, 12,  3,  2,  4,  6 );
	test( 12, 12,  4,  2,  3,  6 );
	test( 12, 12,  3,  3,  4,  4 );
	test( 12, 12,  5,  2,  2+((comm.rank()+1) % 5==0?2:0),  6 );
	test( 12, 12, 11,  1,  1+((comm.rank()+1) %11==0?1:0), 12 );
	test( 12, 12,  4,  3,  3,  4 );
}

TEST_CASE( "structMeshInitOffset: Determining initial offsets" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	Array2i meshSizeGlobal {17,13};
	Array2i nParts { structMeshNParts( meshSizeGlobal, comm.size() ) };
	/* these sizes ignore the fact that at the top and right boundary, the
	 * remainder of the division are added on top. That is, because the offsets
	 * only depend on the mesh sizes which don't lie at those boundaries */
	Array2i meshSize { meshSizeGlobal / nParts };

	for (int j{0}; j<nParts.y(); ++j)
		for (int i{0}; i<nParts.x(); ++i)
			if ( comm.rank() == j*nParts.x() + i )
				REQUIRE( (structMeshInitOffset(meshSizeGlobal, nParts, comm)
					== Array2i{ meshSize.x()*i, meshSize.y()*j }
				).all() );
}

TEST_CASE(
	"equalPartialSums: Calculating segment lenghts with similar partial sums"
){
	SECTION( "Equal loads" ){
		Index nLoads {13};
		ArrayXs loads (nLoads);
		loads = 0.1;
		ArrayXi segments;

		segments.resize(1);
		segments = nLoads;
		REQUIRE( (equalPartialSums(loads, 1) == segments).all() );

		segments.resize(2);
		segments << 6, 7;
		REQUIRE( (equalPartialSums(loads, 2) == segments).all() );

		segments.resize(3);
		segments << 4, 4, 5;
		REQUIRE( (equalPartialSums(loads, 3) == segments).all() );

		segments.resize(4);
		segments << 3, 3, 3, 4;
		REQUIRE( (equalPartialSums(loads, 4) == segments).all() );

		segments.resize(5);
		segments << 3, 3, 3, 3, 1;
		REQUIRE( (equalPartialSums(loads, 5) == segments).all() );

		segments.resize(6);
		segments << 2, 2, 2, 2, 2, 3;
		REQUIRE( (equalPartialSums(loads, 6) == segments).all() );
	}
	SECTION( "Varying loads" ){
		Index nLoads {13};
		ArrayXs loads (nLoads);
		loads << 12, 6, 6, 4, 4, 4, 3, 3, 3, 3, 5, 4, 4;

		ArrayXi segments (5);
		segments << 1, 2, 3, 4, 3;

		REQUIRE( (equalPartialSums(loads,5) == segments).all() );

		/* example from load_balancing.ods */
		ArrayXs horizontalSums (12);
		ArrayXs verticalSums (16);
		segments.resize(3);
		
		segments << 4, 3, 5;
		horizontalSums << 7.2, 8.0, 8.8, 10.4, 10.4, 9.6, 8.8, 7.2, 6.4, 5.6, 5.6, 5.6;
		REQUIRE( (equalPartialSums(horizontalSums, 3) == segments).all() );
		
		horizontalSums << 8, 8, 8, 8, 8.8, 8.8, 8.8, 8.8, 5.6, 5.6, 5.6, 5.6;
		REQUIRE( (equalPartialSums(horizontalSums, 3) == segments).all() );
		

		segments.resize(4);

		segments << 7, 3, 2, 4;
		verticalSums << 0.8, 0.8, 0.8, 1.6, 1.6, 1.6, 2.4, 2.4, 2.4, 3.2, 3.2, 4, 3.2, 2.4, 2.4, 1.6;
		REQUIRE( (equalPartialSums(verticalSums, 4) == segments).all() );
		
		segments << 4, 2, 2, 8;
		verticalSums << 0.6, 0.6, 2.2, 3, 3, 3, 3, 3, 3, 3, 1.4, 0.6, 0.6, 0.6, 0.6, 0.6;
		REQUIRE( (equalPartialSums(verticalSums, 4) == segments).all() );
		
		segments << 6, 2, 2, 6;
		verticalSums <<	1, 1, 1, 1, 1.8, 2.6, 2.6, 4.2, 5, 3.4, 1.8, 1, 1, 1, 1, 1;
		REQUIRE( (equalPartialSums(verticalSums, 4) == segments).all() );
		
		segments << 7, 3, 3, 3;
		verticalSums << 0.8, 0.8, 0.8, 0.8, 1.6, 1.6, 1.6, 1.6, 3.2, 3.2, 3.2, 3.2, 2.4, 2.4, 2.4, 2.4;
		REQUIRE( (equalPartialSums(verticalSums, 4) == segments).all() );
		
		segments << 5, 2, 3, 6;
		verticalSums << 1.2, 1.2, 1.2, 1.2, 3, 3, 3, 3, 1.8, 1.8, 1.8, 1.8, 0.6, 0.6, 0.6, 0.6;
		REQUIRE( (equalPartialSums(verticalSums, 4) == segments).all() );

		segments << 5, 3, 3, 5;
		verticalSums << 1.2, 1.2, 1.2, 1.2, 2.6, 2.6, 2.6, 2.6, 3, 3, 3, 3, 1, 1, 1, 1;
		REQUIRE( (equalPartialSums(verticalSums, 4) == segments).all() );
	}
}

TEST_CASE( "meshPartSizesToOffsets: Calculating offsets from partition sizes" ){
	Array2Xi meshSizes (2, 6);
	meshSizes <<
		1, 2, 3, 2, 3, 1,
		7, 7, 7, 8, 8, 8;
	
	Array2i nParts { 3, 2 };

	Array2Xi offsets;
	offsets.resizeLike(meshSizes);
	offsets <<
		0, 1, 3, 0, 2, 5,
		0, 0, 0, 7, 7, 7;
	
	REQUIRE( (meshPartSizesToOffsets(nParts, meshSizes) == offsets).all() );
}

TEST_CASE( "balancedStructMeshParts: balancing sizes" ){
	/* example from load_balancing.ods */
	Array2i meshSizeGlobal {16,12};
	Array2i nParts {4,3};

	auto loadPerCell = [&](){
		ArrayXs load ( meshSizeGlobal.prod() );
		load = 0.2;
		auto loadReshaped =
			load.reshaped( meshSizeGlobal.x(), meshSizeGlobal.y() );
		
		loadReshaped( seqN(11, 5),  0 ) = 1;
		loadReshaped( seqN( 9, 6),  1 ) = 1;
		loadReshaped( seqN( 6, 7),  2 ) = 1;
		loadReshaped( seqN( 3, 9),  3 ) = 1;
		loadReshaped( seqN( 2, 9),  4 ) = 1;
		loadReshaped( seqN( 2, 8),  5 ) = 1;
		loadReshaped( seqN( 3, 7),  6 ) = 1;
		loadReshaped( seqN( 4, 5),  7 ) = 1;
		loadReshaped( seqN( 5, 4),  8 ) = 1;
		loadReshaped( seqN( 7, 3),  9 ) = 1;
		loadReshaped( seqN( 7, 3), 10 ) = 1;
		loadReshaped( seqN( 8, 3), 11 ) = 1;

		return load;
	};

	auto meshSizesLPC = [](){
		Array2Xi meshSizes (2, 12);
		meshSizes <<
			7, 3, 2, 4,   4, 2, 2, 8,   6, 2, 2, 6,
			4, 4, 4, 4,   3, 3, 3, 3,   5, 5, 5, 5;
		return meshSizes;
	};

	SECTION( "Per cell, local" ){
		REQUIRE( (
			balancedStructMeshParts( meshSizeGlobal, nParts, loadPerCell() )
				== meshSizesLPC()
		).all() );
	}
	SECTION( "Per cell, distributed" ){
		StructMeshPart<UniMesh> part { MPI_COMM_WORLD };
		/* this test is specifically designed for 12 nodes */
		if ( part.comm().size() == 12 ){
			UniMesh mesh;

			if ( part.comm().isMaster() )
				mesh = UniMesh{ meshSizeGlobal, Array2s{1,1} };
			
			part.assignAndDistribute( std::move(mesh) );

			ArrayXs loadGlobal { loadPerCell() };
			ArrayXs load { loadGlobal.reshaped(
				meshSizeGlobal.x(), meshSizeGlobal.y()
			).block(
				part.offset().x(), part.offset().y(),
				part.meshSize().x(), part.meshSize().y()
			).reshaped().eval() };

			Array2Xi meshSizes { balancedStructMeshParts(part, load) };
			MPI_Bcast_H( meshSizes, part.comm(), part.comm().master(), true );
			REQUIRE( ( meshSizes == meshSizesLPC() ).all() );
		}
	}
	SECTION( "Per node" ){
		Array2Xi initialMeshSizes (2, 12);
		initialMeshSizes = 4;
		ArrayXs loadPerNode (12);
		loadPerNode <<
			0.2, 0.4, 0.8, 0.6, 0.4, 1.0, 0.6, 0.2, 0.2, 0.4, 0.6, 0.2;
		Array2Xi meshSizes;
		meshSizes.resizeLike(initialMeshSizes);
		meshSizes <<
			7, 3, 3, 3,   5, 2, 3, 6,   5, 3, 3, 5,
			4, 4, 4, 4,   3, 3, 3, 3,   5, 5, 5, 5;
		REQUIRE( ( balancedStructMeshParts(
				meshSizeGlobal, nParts, initialMeshSizes, loadPerNode
			) == meshSizes ).all()
		);
	}
}

TEST_CASE( "De- and recomposing array" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };
	Array2i globalMeshSize {10, 20};
	Array2i nParts { structMeshNParts( globalMeshSize, comm.size() ) };
	Array2i localMeshSize {
		structMeshPartSize( globalMeshSize, nParts, comm )
	};
	Array2i localOffset {
		structMeshInitOffset( globalMeshSize, nParts, comm )
	};

	SECTION( "Recomposing" ){
		auto test = [&](int nComponents){
			ArrayXXi localArray ( nComponents, localMeshSize.prod() );
			localArray = comm.rank();

			ArrayXXi globalArray { recomposeArray(
				localArray, globalMeshSize, localMeshSize, localOffset
			)};

			Array2Xi allSizes, allOffsets;
			if ( comm.isMaster() ){
				allSizes.resize( 2, nParts.prod() );
				allOffsets.resizeLike(allSizes);
			}
			MPI_Gather(
				localMeshSize.data(), 2, MPI_INDEX,
				allSizes.data(), 2, MPI_INDEX, comm.master(), comm
			);
			MPI_Gather(
				localOffset.data(), 2, MPI_INDEX,
				allOffsets.data(), 2, MPI_INDEX, comm.master(), comm
			);

			if ( comm.isMaster() )
				for (int i{0}; i<comm.size(); ++i)
					REQUIRE( (
						globalArray.reshaped(
							globalMeshSize.x() * nComponents,
							globalMeshSize.y()
						).block(
							allOffsets(0, i)*nComponents, allOffsets(1, i),
							allSizes(0, i)*nComponents, allSizes(1, i)
						) == i
					).all() );
		};
		for( int i{1}; i<5; ++i ) test(i);
	}
	SECTION( "De- and recomposing" ){
		auto test = [&]( int nComponents ){
			ArrayXXi globalArray {
				ArrayXXi::Random( nComponents, globalMeshSize.prod() )
			};

			ArrayXXi localArray { decomposeArray(
				globalArray, globalMeshSize, localMeshSize, localOffset, comm
			)};
			ArrayXXi globalRecomp { recomposeArray(
				localArray, globalMeshSize, localMeshSize, localOffset
			)};

			if ( comm.isMaster() )
				REQUIRE( (globalArray == globalRecomp).all() );
		};
		for( int i{1}; i<5; ++i ) test(i);
	}
	SECTION( "De- and recomposing with Eigen::Ref types" ){
		/* test with different offsets from the start */
		auto test = [&]( int nComponents, int colOffset ){
			ArrayXXi globalArray { ArrayXXi::Random( nComponents,
				globalMeshSize.x()*(globalMeshSize.y()+1) +
				globalMeshSize.y()*(globalMeshSize.x()+1)
			)};

			Array2i newGlobalSize {globalMeshSize.x(), globalMeshSize.y() + 1};
			Array2i newLocalSize {localMeshSize.x(), localMeshSize.y() + 1};

			ArrayXXi localArray { decomposeArray(
				globalArray.block(
					0, colOffset, nComponents,
					globalMeshSize.x()*(globalMeshSize.y()+1)
				), newGlobalSize, newLocalSize, localOffset, comm
			)};

			ArrayXXi globalRecomp { recomposeArray(
				localArray, newGlobalSize, newLocalSize, localOffset
			)};

			if ( comm.isMaster() )
				REQUIRE( (
					globalArray.block(
						0, colOffset, nComponents,
						globalMeshSize.x()*(globalMeshSize.y()+1)
					) == globalRecomp
				).all() );
		};
		for( int i{1}; i<5; ++i )
			for (int j{0}; j<4; ++j)
				test(i,j);
	}
	SECTION( "De- and recomposing using mesh partitions" ){
		auto test = [&]( int nComponents, int nGhostL ){
			StructMeshPart<UniMesh> part { comm };
			UniMesh mesh;

			if ( part.comm().isMaster() )
				try{
					mesh = UniMesh{ globalMeshSize, Array2s{1,1}, {0,0},
						std::max(1, nGhostL)
						// 0
					};
				} catch( const std::exception& e ){
					std::cerr << e.what() << "\n";
					std::abort();
				}

			part.assignAndDistribute( std::move(mesh) );


			ArrayXXi globalArray { ArrayXXi::Constant(
				nComponents, (globalMeshSize + 2*nGhostL).prod(), 1
			) };

			ArrayXXi localArray { decomposeArray(globalArray, part, nGhostL) };

			ArrayXXi globRecomp { recomposeArray(localArray,  part, nGhostL) };

			if ( part.comm().isMaster() )
				REQUIRE( (globalArray == globRecomp).all() );
		};
		for( int j {0}; j < 4; ++j )
			for( int i {1}; i < 4; ++i)
				test(i, j);
	}
}